go Context 设计与实现

Posted on

版本：go 1.19

在前一篇文章中我们讨论了 go Context 的一些常见使用方式，今天我们再来从源码的角度深入了解一下 Context 的设计与实现。 Context 的源码数量不多，去掉注释大概只有两三百行，但是包含的信息量巨大（所以本文也比较长），而且设计得非常巧妙，值得读一读。 然后，下面的 图解 propagateCancel 这一小节的几个图描述了 Context 的工作机制，如果不想看代码，可以直接拉到下面。

再了解一下 chan

在开始本文之前，先来了解一下 Context 实现的关键：chan，对于 chan（再准确一点，我们这里讨论的其实是只读 chan），我们需要清楚以下几点：

  • <-ch 表示从 chan 中获取值。
  • <-ch 在通道（chan）尚未关闭的时候，会一直阻塞，直到通道接收到值。所以有时候通过 select 语句来监听 chan，从而实现协程间的通信。
  • <-ch 在通道（chan）关闭了之后，会立即返回，但是返回的是 chan 关联类型的零值，如果我们还需要判断是否是因为关闭才返回的话，需要用两个值来接收 <-ch 的返回值，如 v, ok := <-ch，ok 表明了通道是否已经关闭，如果是关闭导致它返回，则返回的是 false。

下面这个例子展示了 Context 实现的关键（通过 close(chan)，所有 <-chan 会返回，本质上来说是一种 "广播机制"）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
package main

import (
	"fmt"
	"time"
)

func main() {
	// 创建一个 chan，类型是 struct{}
	ch := make(chan struct{})

	go func() {
		select {
		// 这个 case 会在 chan 关闭或者收到值的时候执行，
		// 在这里的情况是关闭了 chan。
		case v, ok := <-ch:
			if !ok {
				// 输出 "chan ch is closed."
				fmt.Println("chan ch is closed.")
			}
			// 关闭 chan 之后得到的是 ch 的零值，也就是一个空结构体实例
			fmt.Println(v) // {}
		}
	}()

	// 关闭 chan，所有从 chan 读取的操作都会立即返回。
	// 关闭 chan 之后，<-ch 返回的第一个值是 chan 对应类型的零值，第二个参数是 false。
	// 如果不是关闭的 chan，第二个参数是 true，表示可以从 chan 获取到数据。
	close(ch)

	// 防止程序退出看不到效果
	time.Sleep(time.Second)

	// {}，chan struct{} 关闭后，从中获取值的时候会立即返回一个空结构体实例
	fmt.Println(<-ch)
}

在 Context 中，context.Done() 方法返回的 chan，不会接收任何的值，但是在调用 CancelFunc 的时候，会关闭这个 chan，因此 所有的 <-context.Done() 会返回一个零值，返回什么不重要，重要的是，它返回的时候就代表着被上游取消了（代表的是一个取消信号）。

Context 的 UML 图

先来看看它的 UML 图，在后面会陆续展开细说。

  • 首先，有两个关键的接口，分别是 Context 和 canceler。Context 里面的接口是我们实际开发的时候用的，而 canceler 是 context 包内部使用， 这个 canceler 接口定义了一个 cancel 方法，这个方法就是用来发送取消信号的。
  • emptyCtx 代表一个空的 Context，往往用作根 Context，valueCtx 在父 Context 的基础上加了一个键值对。
  • cancelCtx 同时实现了 Context 和 canceler 接口，表示一个可以取消的 Context。
  • timerCtx 在 cancelCtx 的基础上加了一个定时器，表示可以在指定时间之后由定时器进行取消操作。又或者由开发者自行取消。

context 包的结构体、方法说明

  • Context 接口：定义了 Context 接口的四个方法
  • emptyCtx 结构体：一个空 Context
  • CancelFunc 函数类型：Context 的取消函数
  • canceler 接口：Context 取消接口
  • cancelCtx 结构体：实现了取消接口的 Context
  • timerCtx 结构体：超时会取消的 Context
  • valueCtx 结构体：可以存储键值对的 Context
  • Background() 函数：返回空 Context，常作为根 Context
  • TODO() 函数：返回一个空 Context，在需要 Context 的地方又没有合适的 Context 就用这个
  • WithCancel() 函数：基于父 Context，创建一个可取消的 Context
  • newCancelCtx() 函数：创建一个可取消的 Context
  • propagateCancel() 函数：将节点挂载到上游第一个 cancelCtx 上，又或者启动协程监听 Context 取消事件
  • parentCancelCtx() 函数：返回上游的第一个 cancelCtx
  • removeChild() 函数：移除 Context 节点
  • init() 函数：包初始化函数，创建了一个关闭的 chan
  • WithDeadline() 函数：创建一个有 deadline 的 Context
  • WithTimeout() 函数：创建一个有 timeout 的 Context
  • WithValue() 函数：创建一个存储键值对的 Context

Context 接口

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
type Context interface {
    // 返回一个 channel，当 context 被取消或者到了 deadline 的时候，
	// 这个 channel 会被 close，从而 <-chan struct{} 会返回。
	// 在没有关闭之前，一直阻塞，因为不会有任何地方往这个 channel 中发送值。
    Done() <-chan struct{}

    // 在 channel Done 返回的 channel 关闭后，返回 context 取消原因。
    Err() error

    // 返回 context 是否会被取消以及自动取消时间（即 deadline）
	// ok 为 true，表明设置了 deadline，第一个返回值就是设置的 deadline
	// ok 为 false，表示没有设置 deadline，第一个返回值没意义。
    Deadline() (deadline time.Time, ok bool)

    // 获取 key 对应的 value
    Value(key interface{}) interface{}
}

Context 接口定义了 4 个方法，它们都是幂等的，也就是说连续多次调用同一个方法，得到的结果都是相同的。

Deadline() 示例：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
package main

import (
	"context"
	"github.com/davecgh/go-spew/spew"
	"time"
)

func main() {
	ctx, _ := context.WithTimeout(context.Background(), time.Second)
	// 输出 ctx 的 deadline，具体时间为 1 秒之后
	spew.Dump(ctx.Deadline())

	ctx1 := context.Background()
	// ctx1 的超时时间是一个零值
	spew.Dump(ctx1.Deadline())
	
	// 输出：
	// (time.Time) 2022-11-19 11:45:38.702281 +0800 CST m=+1.000233039
	// (bool) true
	// (time.Time) 0001-01-01 00:00:00 +0000 UTC
	// (bool) false
}

我们通过 Deadline() 方法可以知道当前拿到的 Context 参数是否设置了 deadline，以及 deadline 是什么时候， 从而决定接下来是否还需要做一些操作，如果时间太少的话，就可以考虑不做了，因为最终的结果还是超时。

canceler 接口

先看源码：

1
2
3
4
type canceler interface {
    cancel(removeFromParent bool, err error)
    Done() <-chan struct{}
}

实现了 canceler 两个方法的 Context，就表明该 Context 是可取消的。

cancel 方法的第一个参数 removeFromParent 表示的是，是否从父 Context 移除自身，这是因为 Context 是一个树状结构。 在 Context 取消的时候，它会给所有派生的 Context 也发送取消信号，所以派生新的 Context 的时候会记录从当前 Context 派生出去的 Context。

但同样的，在 Context 被取消的时候，父 Context 也就再也不需要给这个 Context 发送取消信号啥的。

我们可以看这个图，Context 派生出了三个 Context，当 child 3 这个 Context cancel 的时候，只会影响到 child 3-1 和 child 3-2 以及其自身， cancel 之后，根结点的 Context 再发送取消信号，child 3 就再也收不到了，因为它已经从这棵树中移除。

emptyCtx 结构体

emptyCtx 本身没有什么实际的作用，一般用作根 Context，比如在 main 函数里面创建的 Context。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
type emptyCtx int

func (*emptyCtx) Deadline() (deadline time.Time, ok bool) {
	return
}

func (*emptyCtx) Done() <-chan struct{} {
	return nil
}

func (*emptyCtx) Err() error {
	return nil
}

func (*emptyCtx) Value(key any) any {
	return nil
}

func (e *emptyCtx) String() string {
	switch e {
	case background:
		return "context.Background"
	case todo:
		return "context.TODO"
	}
	return "unknown empty Context"
}

我们使用 context.Background() 或者 context.TODO() 的时候返回的就是一个 emptyCtx：

1
2
3
4
5
6
7
8
9
10
11
12
var (
	background = new(emptyCtx)
	todo       = new(emptyCtx)
)

func Background() Context {
    return background
}

func TODO() Context {
    return todo
}

emptyCtx 永远不会被取消，也没有值和 deadline。TODO 用在需要 Context 但又没有合适的 Context 可以用的时候。

cancelCtx 结构体

cancel 的操作实际上只会做一次，后续调用 cancel 的时候会返回第一次 cancel 的结果，cancel 是一个幂等操作。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
// 可以被取消，取消的时候，所有实现了 canceler 接口的派生出来的 Context 也会被取消。
type cancelCtx struct {
	// cancelCtx 也实现了 `Context` 接口
	Context

	// mu 用以保护后面的 done、children、err 字段
	mu       sync.Mutex
	
	// 是一个 chan struct{}，懒汉式创建，
	// 在第一次 cancel 的时候被关闭
	done     atomic.Value
	// 记录所有可以取消的子 Context
	// 在第一次 cancel 的时候会被设置为 nil。
	children map[canceler]struct{}
	// 在第一次 cancel 的时候会被设置为非 nil 的值
	err      error
}

cancelCtx.done 是一个支持原子操作的 chan struct{}。

先来看看它的 Done() 方法实现：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
// 返回一个只读的 chan，但没有任何地方会往这个 chan 写入数据，
// cancel 的时候会关闭这个 chan，从而任何 <-ch 的操作都会立即返回。
func (c *cancelCtx) Done() <-chan struct{} {
	// 如果 done 这个 chan 已经初始化了，就直接返回。
	d := c.done.Load()
	if d != nil {
		return d.(chan struct{})
	}
	// 如果 done 还没初始化，则会进行初始化。
	// 也就是上面说的 "懒汉式" 的创建方式，只有在需要的时候才会初始化。
	c.mu.Lock()
	defer c.mu.Unlock()
	d = c.done.Load()
	if d == nil {
		d = make(chan struct{})
		c.done.Store(d)
	}
	return d.(chan struct{})
}

再来看看 Err 方法：

1
2
3
4
5
6
7
func (c *cancelCtx) Err() error {
	// 使用 mu 保证并发安全，本质是 return c.err
	c.mu.Lock()
	err := c.err
	c.mu.Unlock()
	return err
}

cancel 方法

然后再来看 cancel 方法的实现：

cancel 方法做了如下操作：

  • 关闭 c.done
  • 取消 c 的所有孩子 Context
  • 如果 removeFromParent 为 true，会将 c 从其父 Context 的 children 属性中移除
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
// 发送取消信号
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
	// 必须传递一个 err
	if err == nil {
		panic("context: internal error: missing cancel error")
	}
	// 如果已经取消，直接返回。（幂等的设计）
	c.mu.Lock()
	if c.err != nil {
		c.mu.Unlock()
		return // already canceled
	}
	// 记录取消原因，在调用 c.Err() 的时候会返回这个原因
	c.err = err
	// 关闭 done 这个通道，通知其他协程
	d, _ := c.done.Load().(chan struct{})
	if d == nil {
		c.done.Store(closedchan)
	} else {
		close(d)
	}
	// 遍历它的所有子结点，并对其子结点进行取消操作
	for child := range c.children {
		child.cancel(false, err)
	}
	// 将子结点置空
	c.children = nil
	c.mu.Unlock()

	// 从父结点中移除自己
	if removeFromParent {
		removeChild(c.Context, c)
	}
}

WithCancel 方法

我们再来看看创建 cancelCtx 的方式：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
// 返回值里的 Context 的 Done 方法返回的 channel 关闭或者 parent 被 cancel 的时候，
// 返回值的 CancelFunc 会被执行。
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	// 必须从其他 Context 派生
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	c := newCancelCtx(parent)
	// 将 c 挂靠到 parent 的 children 属性中，
	// 从而在 parent 取消的时候，可以感知得到。
	propagateCancel(parent, &c) // 具体实现后面有详细说明 
	return &c, func() { c.cancel(true, Canceled) }
}

// 创建一个 cancelCtx 实例
func newCancelCtx(parent Context) cancelCtx {
    return cancelCtx{Context: parent}
}

我们看到这里的 CancelFunc 的里面调用 cancel 的时候，第一个参数是 true，这表示在取消的时候，需要从 parent 中移除自身。

parentCancelCtx 方法

parentCancelCtx 方法用以从 parent 开始直到根节点的路径搜索第一个 cancelCtx，会跳过中间的 valueCtx。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
// 返回当前节点到根节点路径上的第一个 *cancelCtx。
// 如果 parent 就是 *cancelCtx，那么返回的就是 parent。
// 如果不是，它会从当前结点往 Context 树根结点遍历，找到父结点中的第一个 *cancelCtx，假设是 p。
// 然后检查 p.done 是否跟 parent.Done() 一样，不一样的话意味着 *cancelCtx 已经被包装在自定义实现中了，
// 这个时候，我们不应该绕过它，直接返回 nil 和 false。
// （注意：如果是我们的结构体嵌套了 Context，那么一样会被当做普通的 Context 处理。）
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
	// 判断 parent 的 done 是否已经关闭或者并没有 done chan。
	// 如果是，则返回 nil 和 false
    done := parent.Done()
	// done 为 nil 表示 parent 或者到根 Context 这条路径上并没有 *cancelCtx（只有 valueCtx 或 emptyCtx）。
    if done == closedchan || done == nil {
        return nil, false
    }
	// 判断 parent 是否是一个 cancelCtx
	// 如果不是，则返回 nil 和 false
	// 讲道理，parent.Value(&cancelCtxKey) 的返回值只有两种情况：
	// emptyCtx（找到根节点也没找到）或者 *cancelCtx（找到了）
	// （parent.Value 实现细节见下面的 value 那一小节）
    p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
    if !ok {
        return nil, false
    }
	// 执行到这里的时候：p 是一个 *cancelCtx
	// 判断 parent.Done() 和 p.done 是否相等：
	// 不等则意味着 *cancelCtx 已经被包装在自定义实现中了，这个时候，我们不应该绕过它。
	// 详细请参考：go issue 28728（google）
    pdone, _ := p.done.Load().(chan struct{})
    if pdone != done {
        return nil, false
    }
    return p, true
}

// 从父结点移除自己（从 parent 移除 child）
func removeChild(parent Context, child canceler) {
	p, ok := parentCancelCtx(parent)
	if !ok {
		return
	}
	// 从父结点的 children 中移除 child
	p.mu.Lock()
	if p.children != nil {
		delete(p.children, child)
	}
	p.mu.Unlock()
}

propagateCancel 函数

这个函数会在两个地方调用，一个是 WithCancel，另一个是 WithDeadline，它的主要作用是，找到 parent 以及其父级 Context 路径上 第一个 cancelCtx，目的是，将 child 挂载到找到的这个 cancelCtx 的 children 属性上，从而在这个 cancelCtx 取消的时候， 可以通过遍历 cancelCtx.children 对 child 进行通知。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
// 由 parent 往根节点搜索第一个 cancelCtx，如果找到则将 child 写入到 cancelCtx.children 中。
// 如果找到的 cancelCtx 自定义了 Done，则启动协程监听 cancelCtx.Done()。 
func propagateCancel(parent Context, child canceler) {
	// 如果 Context 树上完全不存在 cancelCtx，则直接返回
	done := parent.Done()
	if done == nil {
		return // parent is never canceled
	}

	// 如果 parent 已经取消，则直接取消 child
	select {
	case <-done:
		// parent is already canceled
		child.cancel(false, parent.Err())
		return
	default:
	}

	// 往根节点搜索第一个 cancelCtx
	if p, ok := parentCancelCtx(parent); ok {
		// 找到了，但是已经取消了，则取消 child
		p.mu.Lock()
		if p.err != nil {
			// parent has already been canceled
			child.cancel(false, p.err)
		} else {
			// 找到了，尚未取消。
			// 将 child 写入到 p 的 children 属性中。
			// p.children 是懒汉式创建的。
			if p.children == nil {
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}
		}
		p.mu.Unlock()
	} else {
		// 执行到这里的原因是：
		// 用户自定义了 Done 通道（跟 parent 不是同一个 done），
		// 所以不能以父节点路径上的 done 来决定 child 是否取消，
		// 需要通过启动新协程的方式来监听 Done 通道，从而可以正常取消 parent 的孩子节点。
		atomic.AddInt32(&goroutines, +1)
		go func() {
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err())
			case <-child.Done():
			}
		}()
	}
}

图解 propagateCancel

propagateCancel 的实现可能只看代码不好理解，所以我画了几个图来帮助我们理解：

我们可以看到 propagateCancel 里面有一个 parentCancelCtx，对于 parentCancelCtx，下一小节有比较详细的说明。 这个图描述了 parentCancelCtx 的实际执行过程，在我们调用 propagateCancel 的时候，搜索 cancelCtx 的过程：

  • 首先，我们知道，Context 是一个树状结构，每一个 Context 都可以派生出子 Context。
  • 图中 parentA 是 child 到根节点 emptyCtx 路径上的第一个 cancelCtx。
  • parentCancelCtx 拿到的是 parentA，然后将 child 写入到 parentA 的 children 属性中。
  • 从而在我们手动取消 parentA 的时候，parentA 可以通过遍历 children 的方式，告知 child 协程取消了。

注意：搜索的时候会跳过 valueCtx。

在 child 到 emptyCtx 路径上搜索第一个 cancelCtx 的过程：

我们取消的过程大概如下图：

我们理解取消的过程的时候，可以忽略掉 Context 树中那些非 cancelCtx 节点，正如上图这样，实际上取消过程只涉及到 parentA 以及 child， 其他节点如何并不影响。

如果觉得这个图不太好懂，可以再看看这个图：

用户覆盖了 done 的特殊情况：

如果用户覆盖了 done 通道，这表明用户想自行控制什么时候 parentB 结束。（也就是说，parentB 脱离了路径上 cancelCtx 的控制， 也就是假设 parentA 还是 cancelCtx，在 parentA 取消的时候，parentB 是收不到信号的，parentB 收到信号是在其 Done() 返回的通道关闭的时候。）

这个时候因为我们从 parentB 派生出了一个新的 cancelCtx（child），所以 parentB 需要对 child 进行控制， 也就是说在 parentB 取消的时候，也取消 child。这种情况下，就是通过 propagateCancel 里面的协程里面实现的。

parentCancelCtx 函数

parentCancelCtx 的描述比较晦涩，如果没有实际的例子我们很难看得懂它的意思。里面有一个比较， 是针对 parent.Done() 和 p.done.Load().(chan struct{}) 的，源码里面判断如果这两者不一样，则返回 nil 和 false。

注释里说，如果两者不一样，我们不应该绕过它（bypass it），但是这里的绕过是什么意思呢？我们来看一个例子：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
package c

import (
	"testing"
	"time"
)

type A struct {
	Context
	ch <-chan struct{}
}

func (a *A) Done() <-chan struct{} {
	return a.ch
}

func TestCancel(t *testing.T) {
	// 创建一个 cancel context
	ctx, cancel0 := WithCancel(TODO())

	// 创建一个 A 实例
	// 这个实例可以内嵌了 Context，所以可以当作 Context 使用，
	// 但是我们覆盖了 Context 本身的 Done 方法。
	ch := make(<-chan struct{})
	a := A{ctx, ch}

	// 因为我们覆盖了 Done 方法，所以 go 底层会认为是开发者想要
	// 自行控制协程取消，所以在 WithCancel 的时候并不会把 ctx1
	// 挂载到 a 的 children 属性下，这样一来，
	// go 底层只能再启动一个协程来监听 a 的 Done chan，
	// 从而可以在 a cancel 的时候可以正常通知到 ctx1。
	ctx1, cancel := WithCancel(&a)

	go func() {
		time.Sleep(time.Millisecond * 10)
		cancel()
	}()
	
	// ctx2 会写入到 ctx1 的 children 属性中，
	// 这样就不需要启动新的协程来监测 ctx1 的 done。
	ctx2, cancel2 := WithCancel(ctx1)

	time.Sleep(time.Millisecond * 20)
}

下面这个图描述了上面这个例子中的 Context 树结构：

我们调用了三次 WithCancel，这三次的效果都不太一样：

  • 第一次调用的时候，parent.Done 返回 nil，这个时候，我们取消只有调用 cancel0 这一种途径，也就是手动取消。
  • 第二次调用的时候，parent.Done 返回的不是 nil，但是和 go 语言底层的那个 done 属性不一致（一个是 A.Done() ，另一个是 cancelCtx.done）。这种情况下，go 底层就知道，开发者自己定义了一个 done 通道，这个时候，会需要另外启动一个协程来监听 A.Done() 返回的 done，从而可以在 A 结束的时候，通知 A 的孩子 ctx1。
  • 第三次调用的时候，parent.Done 返回的不是 nil，而且和 cancelCtx.done 相等，说明用户没有重写 Done 方法，这样就可以直接将 ctx2 挂载到 ctx1 的 children 属性上，而不用另外启动协程来监听 ctx1 的 done。

具体怎么实现的可以看上面的 propagateCancel 这一小节。

之所以这样是为了给开发者一定的控制权，如果忽略了用户自定义的 Done 方法，那么可能取消的操作用户就无法控制了。 但我们覆盖 Done 方法就是为了可以自主去控制取消的操作。

child 什么时候从父 Context 移除？

如果我们足够细心，就会发现我们在 cancel 的时候，有的地方需要将 child 从 Context 中移除，而有的地方不需要，那什么时候需要呢？

需要移除的情况：

  1. WithCancel 派生出新的 Context 的时候，假设叫 root，这个时候派生的这个 root 也是可以继续派生出新的 Context 的，而这个 root 对于它的子孙 Context 它就是根节点，所以当 root 被取消的时候，它和它的子孙 Context 也要被取消了，所以以 root 为根节点的子树需要被移除。
  2. WithDeadline 里面，当给定的 d 其实已经小于当前时间的时候（也就是父 Context 已经超时了），这个时候会将刚挂载到父节点的 timerCtx 移除，同时返回的 CancelFunc 中，cancel 的第一个参数是 false，因为它已经被移除了。

不需要移除的情况：

  1. 在 propagateCancel 中监测到 parent 已经被取消的时候，因为这个时候 child 并没有关联上 parent，所以自然也没有移除的这种操作。
  2. 就是上面提到的第二种情况中，WithDeadline 的时候就监测到 deadline 已经比当前时间小了（超时了）。
  3. 在 cancelCtx 的 cancel 方法里面，遍历 cancelCtx 的孩子节点的时候，不需要做移除的操作，因为 cancelCtx 本身就需要被从 Context 树中移除。
  4. timerCtx 在没有挂载到 parent 上就已经过期了。

timerCtx 结构体

timerCtx 是一个带有定时器的 cancelCtx，我们既可以手动取消，也可以由底层定时器在到达 deadline 的时候进行取消。

1
2
3
4
5
6
7
8
9
// timerCtx 嵌套了 cancelCtx，这表示我们可以手动取消。
// 另外还有一个定时器，这个定时器的执行时间定在 deadline 这个时间点，
// 一旦时间到了，就会调用 cancelCtx 的 cancel 方法。
type timerCtx struct {
	cancelCtx
	timer *time.Timer // Under cancelCtx.mu.

	deadline time.Time
}

我们有两种方法来创建 timerCtx，一个是 WithDeadline：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
// 基于 parent 创建一个 cancelCtx，内嵌到 timerCtx 中。
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
	// 不能基于 nil 创建一个 timerCtx
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	// 如果当前设置的超时时间比 parent 设置的超时时间更长，
	// 那么不用 timerCtx 开启定时器了，因为 parent 会先到期取消，
	// 这里再启动一个定时器也没有执行的机会了。
	if cur, ok := parent.Deadline(); ok && cur.Before(d) {
		// The current deadline is already sooner than the new one.
		return WithCancel(parent)
	}
	// 创建一个 timerCtx
	c := &timerCtx{
		cancelCtx: newCancelCtx(parent),
		deadline:  d,
	}
	// 将刚创建的 timerCtx 挂载到父 Context 中
	propagateCancel(parent, c)
	// 判断还有多久到达 deadline
	dur := time.Until(d)
	// 如果 deadline 已经过去了，那么直接执行 timerCtx 的 cancel 逻辑,
	// 同时移除跟父节点的关联。（创建了还没来得及启动定时器就到期了）
	if dur <= 0 {
		c.cancel(true, DeadlineExceeded) // deadline has already passed
		// cancel 不再需要从父节点移除自身，上一行已经移除了
		return c, func() { c.cancel(false, Canceled) }
	}
	// 启动一个定时器，在到达 deadline 的时候执行 cancel 操作。
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.err == nil {
		c.timer = time.AfterFunc(dur, func() {
			c.cancel(true, DeadlineExceeded)
		})
	}
	// 除了定时器之外，返回一个 CancelFunc 给用户提供自行取消的方式。
	return c, func() { c.cancel(true, Canceled) }
}

另外一个是 WithTimeout：

WithTimeout 本质上是对 WithDeadline 的调用而已，只不过描述到期时间的方式不一样而已。 WithDeadline 描述的是具体的到期时间，WithTimeout 描述的是多久以后的时间，两者其实都代表未来的某一个时间点。

1
2
3
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
	return WithDeadline(parent, time.Now().Add(timeout))
}

valueCtx

1
2
3
4
5
6
// valueCtx 在父 Context 的基础上，带有一对键值对。
// 实现了一个 Value 方法，其他方法都是调用 parent 的。
type valueCtx struct {
	Context
	key, val any
}

WithValue 方法

WithValue 方法一般用在请求范围内的数据共享，WithValue 方法很简单，就是在 parent 的基础上加上了一个 key 和 一个 value。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
// 返回的子 Context 里面有 key val 对
func WithValue(parent Context, key, val any) Context {
	// parent 不能为 nil
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	// key 不能为 nil
	if key == nil {
		panic("nil key")
	}
	// key 必须是可以比较的，
	// 因为在获取值的时候需要进行 key 的比较。
	if !reflectlite.TypeOf(key).Comparable() {
		panic("key is not comparable")
	}
	// 返回 valueCtx
	return &valueCtx{parent, key, val}
}

valueCtx

valueCtx 是从其他 Context 派生出来的，所以内嵌了 Context 接口，同时还有两个字段是 key 和 val，表示的是父 Context 传递的键值。

1
2
3
4
type valueCtx struct {
	Context
	key, val any
}

value 方法

value 方法用以从 Context 中获取对应的值，它会从 Context 树自底向上进行递归搜索，具体来说会有以下几种情况：

  • 如果 ctx 是 *valueCtx，则会判断 key 是否等于 ctx 里面的 key，如果相等，返回 ctx.val。否则，再去搜索 ctx 的父 Context。
  • 如果 ctx 是 *cancelCtx，同时 key 是 &cancelCtxKey，则会返回 ctx。否则，会继续搜索 ctx 到根结点这个路径上的第一个 cancelCtx。
  • 如果 ctx 是 *timerCtx，同时 key 是 &cancelCtxKey，则会返回 ctx.cancelCtx。否则，会继续搜索 ctx 到根结点这个路径上的第一个 cancelCtx。
  • 如果 ctx 是 *emptyCtx，则会返回 nil。（因为这时候是最顶层的 Context 了，也找不到对应的值）。
  • 如果都不是以上的几种情况，则有可能是开发者自定义的 Context 实现，则直接返回 c.Value(key)。

它要解决的问题是：

  • 获取父级 Context 中 WithValue 共享的值。
  • 获取父级 Context 中最靠近当前节点的 cancelCtx（非常重要：它的一个很重要的作用是，将当前节点设置为这个 cancelCtx 的 children，从而可以实现在这个父级的 cancelCtx 取消的时候，当前的 Context 可以感知到）。
  • 如果是开发者自己实现的 Context，则直接调用用户自定义的 Value 方法。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
// 根据 key 从 c 中获取对应的值，会从 Context 树自底向上递归搜索。
func value(c Context, key any) any {
	for {
		switch ctx := c.(type) {
		case *valueCtx:
			if key == ctx.key {
				return ctx.val
			}
			c = ctx.Context
		case *cancelCtx:
			if key == &cancelCtxKey {
				return c
			}
			c = ctx.Context
		case *timerCtx:
			if key == &cancelCtxKey {
				return &ctx.cancelCtx
			}
			c = ctx.Context
		case *emptyCtx:
			return nil
		default:
			return c.Value(key)
		}
	}
}

看下面的图可能会更加直观：

这个方法大概做的事情是，从当前的 Context 中查找指定的 key，如果找不到则递归地从其父级 Context 上找。

但我们最需要关注的是，它这里面 case *cancelCtx 和 case *timerCtx 的逻辑表明了，在我们调用 parent.Value(&cancelCtxKey) 的 时候，实际上获取到的是当前 Context 到根 Context 上第一个 cancelCtx。

一个更简单的图是下面这样的（当然在实际中，会有多个子节点，这里假设都是只有一个子节点）：

为什么是通过关闭 chan 的方式取消？

上面我们说了，在 Context 取消的时候，是通过关闭 chan 的方式来实现的，那么为什么要这么做呢？这是因为，如果说要通过往 chan 写入数据的方式来 通知其他子孙 Context 的话，我们就需要有多少个子孙 Context 就要往 chan 里面发多少次，但是如果选择使用 close 的方式的话， 我们就完全不用管派生出了多少个可以 cancel 的 Context，因为一旦 chan 关闭了，所有的 <-chan 操作立即得以返回，这同样也实现了通信。 但是 close 这种方式无疑更加高明，更加简洁，当然也更加巧妙。

总结

本文主要讲了如下内容：

  • 介绍了 context 包的几个 Context 结构体，其中 emptyCtx 一般用作根 Context，valueCtx 一般用作请求范围内的数据共享，而 cancelCtx 给开发者控制下游 Context 提供了一种很好的方式，timerCtx 在 cancelCtx 的基础上加了一个定时器，时间到会发送取消信号。
  • timerCtx 有两种取消方式，一种是开发者手动取消，这个和 cancelCtx 一致，另外一种方法是到达 deadline 的时候，由定时器来取消。
  • go 里面 Context 取消的时候，是通过关闭 chan 的方式来让下游的 Context 感知的，因为 chan 的工作机制就是如果被关闭则调用 <-chan 会立即返回。
  • 创建 cancelCtx 的时候，会将派生的 Context 挂载到上游 Context 中第一个 cancelCtx 的 children 上，这样在上游取消的时候，这个派生的 Context 可以感知得到。
  • 上游取消的信号会跨过中间的所有 valueCtx，传达到下游那些 cancelCtx，在取消的时候，Context 会从 Context 树中移除。
  • 如果开发者实现了自己的 Done 通道，并且返回的是跟底层 cancelCtx 中不一样的 done 通道，则会导致 go 底层启动一个协程来监测这个被覆盖的 done 通道。
  • value 也会从 Context 树中自底向上搜索，直到根节点。